import os
import re


def merge_split_lines(file_path):
    """修复因换行导致的错误分割数据"""
    merged_data = []
    current_label = None
    text_parts = []

    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.rstrip('\n')  # 保留原始换行信息
            # 检测有效数据行（数字开头 + 制表符）
            if re.match(r'^\d+\t', line):
                # 保存前一个完整数据
                if current_label is not None:
                    merged_line = f"{current_label}\t{' '.join(text_parts)}"
                    merged_data.append(merged_line + '\n')
                # 开始新记录
                current_label, _, text = line.partition('\t')
                text_parts = [text.strip()]
            else:
                # 延续文本内容
                if current_label is not None and line.strip():
                    text_parts.append(line.strip())

    # 添加最后一个记录
    if current_label is not None:
        merged_data.append(f"{current_label}\t{' '.join(text_parts)}\n")
    return merged_data


def process_tsv_files(base_dir):
    """批量处理目录下的TSV文件"""
    for filename in ['train.tsv', 'dev.tsv', 'test.tsv']:
        input_path = os.path.join(base_dir, filename)
        output_path = os.path.join(base_dir, f"corrected_{filename}")

        cleaned_data = merge_split_lines(input_path)

        with open(output_path, 'w', encoding='utf-8') as f:
            f.writelines(cleaned_data)
        print(f"生成文件: {output_path} (共处理{len(cleaned_data)}条数据)")


if __name__ == "__main__":
    # 根据截图中的目录结构设置路径
    data_dir = os.path.join("thucnews_mini_tsv")
    process_tsv_files(data_dir)